import importlib
import pickle
from pathlib import Path
import pandas as pd
from sklearn.pipeline import Pipeline
from sklearn.model_selection import GridSearchCV, cross_val_score
from sklearn.ensemble import StackingClassifier, StackingRegressor


class ModelManager:
    def __init__(self, config, logger, task_type):
        self.config = config
        self.logger = logger
        self.task_type = task_type
        self.model_dir = Path(config.get("paths.model_dir"))
        self.model_dir.mkdir(exist_ok=True)
        self.random_state = config.get("global_params.random_state")
        self.cv_folds = config.get("global_params.cv_folds")
        self.top_n = config.get("global_params.top_n_models")

    def _load_model_class(self, class_path: str):
        """动态加载模型类（如"sklearn.linear_model.LogisticRegression"）"""
        module_name, class_name = class_path.rsplit(".", 1)
        module = importlib.import_module(module_name)
        return getattr(module, class_name)

    def get_model_configs(self) -> dict:
        """获取当前任务类型对应的模型配置"""
        all_models = self.config.get("models", {})
        task_models = all_models.get(self.task_type, {})
        if not task_models:
            raise ValueError(f"未找到{self.task_type}任务的模型配置")
        return task_models

    def train_base_models(self, X_train, y_train, preprocessor) -> pd.DataFrame:
        """训练基础模型并返回Top N模型"""
        model_configs = self.get_model_configs()
        results = []

        for model_name, cfg in model_configs.items():
            self.logger.info(f"\n====== 训练{model_name} ======")
            try:
                # 初始化模型
                model_class = self._load_model_class(cfg["class"])
                model = model_class(**cfg["init_params"])

                # 构建管道（根据任务类型区分分类器/回归器）
                step_name = "classifier" if self.task_type != "regression" else "regressor"
                pipeline = Pipeline([
                    ('preprocessor', preprocessor),
                    (step_name, model)
                ])

                # 网格搜索
                grid_search = GridSearchCV(
                    estimator=pipeline,
                    param_grid=cfg["param_grid"],
                    cv=self.cv_folds,
                    scoring=self._get_scoring_metric(),
                    n_jobs=-1,
                    verbose=1
                )
                grid_search.fit(X_train, y_train)

                # 评估交叉验证性能
                best_pipeline = grid_search.best_estimator_
                cv_score = grid_search.best_score_
                secondary_metric = self._get_secondary_metric()
                cv_secondary = cross_val_score(
                    best_pipeline, X_train, y_train,
                    cv=self.cv_folds, scoring=secondary_metric
                ).mean()

                # 保存结果
                results.append({
                    "模型名称": model_name,
                    "交叉验证主指标": round(cv_score, 4),
                    "交叉验证次要指标": round(cv_secondary, 4),
                    "最佳参数": grid_search.best_params_,
                    "模型管道": best_pipeline
                })
                self.logger.info(f"{model_name}最佳参数: {grid_search.best_params_}")
                self.logger.info(f"{model_name}交叉验证性能 - 主指标: {cv_score:.4f}, 次要指标: {cv_secondary:.4f}")

            except Exception as e:
                self.logger.error(f"{model_name}训练失败: {str(e)}", exc_info=True)
                continue

        # 排序并选择Top N模型
        results_df = pd.DataFrame(results).sort_values(by="交叉验证主指标", ascending=False)
        top_models = results_df.head(self.top_n)
        self.logger.info(f"\nTop {self.top_n}模型:\n{top_models[['模型名称', '交叉验证主指标']]}")

        # 保存模型
        for _, row in top_models.iterrows():
            model_path = self.model_dir.joinpath(f"{row['模型名称']}_最优模型.pkl")
            with open(model_path, "wb") as f:
                pickle.dump(row["模型管道"], f)
            self.logger.info(f"已保存模型: {model_path}")

        return top_models

    def train_stacking_model(self, X_train, y_train, base_models: pd.DataFrame) -> pd.DataFrame:
        """训练堆叠模型"""
        self.logger.info("\n====== 训练堆叠模型 ======")
        estimators = [(row["模型名称"], row["模型管道"]) for _, row in base_models.iterrows()]

        # 选择元模型
        if self.task_type in ["binary", "multiclass"]:
            meta_model = self._load_model_class("sklearn.linear_model.LogisticRegression")(
                class_weight="balanced", random_state=self.random_state
            )
            stacking_clf = StackingClassifier(
                estimators=estimators,
                final_estimator=meta_model,
                cv=self.cv_folds,
                stack_method="predict_proba" if self.task_type != "regression" else "predict"
            )
        else:  # regression
            meta_model = self._load_model_class("sklearn.linear_model.Ridge")(alpha=1.0)
            stacking_clf = StackingRegressor(
                estimators=estimators,
                final_estimator=meta_model,
                cv=self.cv_folds
            )

        # 训练堆叠模型
        stacking_pipeline = Pipeline([('stacking', stacking_clf)])
        stacking_pipeline.fit(X_train, y_train)

        # 评估性能
        main_metric = self._get_scoring_metric()
        cv_main = cross_val_score(stacking_pipeline, X_train, y_train, cv=self.cv_folds, scoring=main_metric).mean()
        secondary_metric = self._get_secondary_metric()
        cv_secondary = cross_val_score(stacking_pipeline, X_train, y_train, cv=self.cv_folds,
                                       scoring=secondary_metric).mean()

        self.logger.info(f"堆叠模型交叉验证性能 - 主指标: {cv_main:.4f}, 次要指标: {cv_secondary:.4f}")

        # 保存模型
        stacking_path = self.model_dir.joinpath("堆叠模型_最优模型.pkl")
        with open(stacking_path, "wb") as f:
            pickle.dump(stacking_pipeline, f)
        self.logger.info(f"已保存堆叠模型: {stacking_path}")

        # 合并结果
        stacking_result = pd.DataFrame([{
            "模型名称": "堆叠模型",
            "交叉验证主指标": round(cv_main, 4),
            "交叉验证次要指标": round(cv_secondary, 4),
            "模型管道": stacking_pipeline
        }])
        return pd.concat([stacking_result, base_models], ignore_index=True)

    def _get_scoring_metric(self) -> str:
        """根据任务类型返回主评分指标"""
        if self.task_type == "regression":
            return "neg_mean_squared_error"  # 回归用MSE（负值，因为GridSearch默认最大化）
        elif self.task_type == "binary":
            return "roc_auc"
        elif self.task_type == "multiclass":
            return "accuracy"

    def _get_secondary_metric(self) -> str:
        """根据任务类型返回次要评分指标"""
        if self.task_type == "regression":
            return "r2"
        else:
            return "accuracy"